#include <iostream>
#include <cstring>  // For memset
#include <chrono>
#include <cryptopp/aes.h>
#include <cryptopp/gcm.h>
#include <cryptopp/filters.h>
#include <cryptopp/modes.h>
#include <cryptopp/osrng.h>
#include "test.pb.h"

using namespace std;
using namespace CryptoPP;

// 固定大小的缓冲区大小
const size_t BUFFER_SIZE = 4096;

// AES-CBC 加密函数
void encryptWithAES(const uint8_t* plaintext, size_t plaintext_size, const uint8_t* key, const uint8_t* iv, uint8_t* ciphertext, size_t& ciphertext_size) {
    CBC_Mode<AES>::Encryption encryption;
    encryption.SetKeyWithIV(key, AES::DEFAULT_KEYLENGTH, iv);

    // 加密
    ArraySink cs(ciphertext, BUFFER_SIZE + AES::BLOCKSIZE);  // 确保缓冲区足够大
    StreamTransformationFilter stf(encryption, new Redirector(cs));
    stf.Put(plaintext, plaintext_size);
    stf.MessageEnd();
    ciphertext_size = cs.TotalPutLength();

//    std::cout << "AES-CBC Encrypted ciphertext size: " << ciphertext_size << std::endl;
}


// AES-CBC 解密函数
void decryptWithAES(const uint8_t* ciphertext, size_t ciphertext_size, const uint8_t* key, const uint8_t* iv, uint8_t* decrypted, size_t& decrypted_size) {
    CBC_Mode<AES>::Decryption decryption;
    decryption.SetKeyWithIV(key, AES::DEFAULT_KEYLENGTH, iv);

    // 解密
    ArraySink ds(decrypted, BUFFER_SIZE);  // 确保缓冲区足够大
    StreamTransformationFilter stf(decryption, new Redirector(ds));
    stf.Put(ciphertext, ciphertext_size);
    stf.MessageEnd();
    decrypted_size = ds.TotalPutLength();

//    std::cout << "AES-CBC Decrypted data size: " << decrypted_size << std::endl;
}


int main() {
    // 创建 Protobuf 消息
    TestEncrypt message;
    message.set_content("加密解密成功");

    // 计算序列化后的消息大小
    int message_size = message.ByteSizeLong();

    // 检查消息是否适合固定大小的缓冲区
    if (message_size > BUFFER_SIZE) {
        std::cerr << "Message size exceeds buffer size" << std::endl;
        return -1;
    }

    // 使用固定大小的缓冲区
    uint8_t buf[BUFFER_SIZE];
    uint8_t key[AES::DEFAULT_KEYLENGTH];  // AES-128 使用 16 字节密钥
    uint8_t iv[AES::BLOCKSIZE];           // AES 使用 16 字节 IV
    uint8_t ciphertext[BUFFER_SIZE + AES::BLOCKSIZE];
    uint8_t decrypted[BUFFER_SIZE];
    size_t ciphertext_size, decrypted_size;

    // 序列化消息到数组
    if (!message.SerializeToArray(buf, BUFFER_SIZE)) {
        std::cerr << "Failed to serialize Protobuf message" << std::endl;
        return -1;
    }


    // 生成密钥和 IV
    AutoSeededRandomPool prng;
    prng.GenerateBlock(key, sizeof(key));
    prng.GenerateBlock(iv, sizeof(iv));

    // 计算1000万次 AES-CBC 加密耗时
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < 10000000; i++) {
        encryptWithAES(buf, message_size, key, iv, ciphertext, ciphertext_size);
    }
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "AES-CBC 加密 1000万次 耗时: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;

    // 计算1000万次 AES-CBC 解密耗时
    try {
        start = std::chrono::high_resolution_clock::now();
        for (int i = 0; i < 10000000; i++) {
            decryptWithAES(ciphertext, ciphertext_size, key, iv, decrypted, decrypted_size);
        }
        end = std::chrono::high_resolution_clock::now();
        std::cout << "AES-CBC 解密 1000万次 耗时: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;
    } catch (const std::runtime_error& e) {
        std::cerr << e.what() << std::endl;
        return -1;
    }

    // 反序列化消息
    TestEncrypt decryptedMessage;
    if (!decryptedMessage.ParseFromArray(decrypted, decrypted_size)) {
        std::cerr << "Failed to parse Protobuf message" << std::endl;
        return -1;
    }

    // 打印解密后的内容
    std::cout << "解密出来的消息内容是: " << decryptedMessage.content() << std::endl;

    return 0;
}
